package com.sluggard.ai.manage;

import com.sluggard.ai.config.MilvusProperties;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.GetLoadStateResponse;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.collection.GetLoadStateParam;
import io.milvus.param.collection.LoadCollectionParam;
import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.milvus.MilvusVectorStore;
import org.springframework.stereotype.Component;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * 类描述：向量存储管理器，根据知识库ID获取或创建VectorStore
 *
 * @author lizheng
 * @created 2025/4/23 10:20
 */
@Component
@RequiredArgsConstructor
public class VectorStoreManage {

    private static VectorStoreManage INSTANCE;

    private static final ConcurrentMap<Long, VectorStore> map = new ConcurrentHashMap<>();

    private final MilvusProperties milvusProperties;
    private final EmbeddingModel embeddingModel;
    private final MilvusServiceClient milvusClient;


    @PostConstruct
    public void init() {
        INSTANCE = this;
    }

    /**
     * 根据知识库ID获取VectorStore，不存在则创建
     *
     * @param knowledgeBaseId 知识库ID
     * @return VectorStore实例
     */
    public static VectorStore getVectorStore(Long knowledgeBaseId) {
        String collectionName = getCollectionName(knowledgeBaseId);
        LoadCollectionParam loadCollectionParam = LoadCollectionParam.newBuilder()
                .withDatabaseName(INSTANCE.milvusProperties.getDatabaseName())
                .withCollectionName(collectionName)
                .build();

        INSTANCE.milvusClient.loadCollection(loadCollectionParam);
        GetLoadStateParam getLoadStateParam = GetLoadStateParam.newBuilder()
                .withCollectionName(collectionName)
                .build();
        R<GetLoadStateResponse> loadState = INSTANCE.milvusClient.getLoadState(getLoadStateParam);
        if (loadState.getStatus() != 0) {
            throw new RuntimeException("加载数据库失败");
        }
        return map.computeIfAbsent(knowledgeBaseId, (key) -> INSTANCE.createVectorStore(key));
    }

    public static String getCollectionName(Long knowledgeBaseId) {
        return INSTANCE.milvusProperties.getKnowledgeBaseKey() + knowledgeBaseId;
    }

    /**
     * 创建新的VectorStore实例
     *
     * @param knowledgeBaseId 知识库ID
     * @return 新创建的VectorStore实例
     */
    public static VectorStore createVectorStore(Long knowledgeBaseId) {
        String collectionName = INSTANCE.milvusProperties.getKnowledgeBaseKey() + knowledgeBaseId;
        return MilvusVectorStore.builder(INSTANCE.milvusClient, INSTANCE.embeddingModel)
                .collectionName(collectionName)
                .databaseName(INSTANCE.milvusProperties.getDatabaseName())
                .indexType(IndexType.IVF_FLAT)
                .metricType(MetricType.L2)
                .batchingStrategy(new TokenCountBatchingStrategy())
                .initializeSchema(true)
                .build();
    }

    /**
     * 移除指定知识库ID的VectorStore
     *
     * @param knowledgeBaseId 知识库ID
     * @return 被移除的VectorStore，如果不存在返回null
     */
    public static VectorStore removeVectorStore(Long knowledgeBaseId) {
        return map.remove(knowledgeBaseId);
    }

    /**
     * 检查指定知识库ID的VectorStore是否存在
     *
     * @param knowledgeBaseId 知识库ID
     * @return 存在返回true，否则返回false
     */
    public static boolean containsVectorStore(Long knowledgeBaseId) {
        return map.containsKey(knowledgeBaseId);
    }

}
